from mcp.server.fastmcp import FastMCP

# Create an MCP server for math operations
mcp = FastMCP("Math Operation Server")

@mcp.tool()
def add(a: float, b: float) -> float:
    """Add two numbers"""
    try:
        return float(a) + float(b)
    except (TypeError, ValueError) as e:
        raise ValueError(f"Invalid input: {e}")

@mcp.tool()
def subtract(a: float, b: float) -> float:
    """Subtract b from a"""
    return a - b

@mcp.tool()
def multiply(a: float, b: float) -> float:
    """Multiply two numbers"""
    try:
        return float(a) * float(b)
    except (TypeError, ValueError) as e:
        raise ValueError(f"Invalid input: {e}")

@mcp.tool()
def divide(a: float, b: float) -> float:
    """Divide a by b"""
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b

@mcp.tool()
def power(base: float, exponent: float) -> float:
    """Calculate base raised to the power of exponent"""
    return base ** exponent

@mcp.tool()
def square_root(x: float) -> float:
    """Calculate square root of x"""
    if x < 0:
        raise ValueError("Cannot calculate square root of negative number")
    return x ** 0.5

if __name__ == "__main__":
    # Add performance monitoring
    import time
    
    def log_performance(func):
        def wrapper(*args, **kwargs):
            start_time = time.time()
            try:
                result = func(*args, **kwargs)
                duration = time.time() - start_time
                print(f"{func.__name__} executed in {duration:.4f} seconds")
                return result
            except Exception as e:
                print(f"Error in {func.__name__}: {str(e)}")
                raise
        return wrapper
    
    # Decorate all tools with performance monitoring
    for tool in mcp._tools.values():
        tool.func = log_performance(tool.func)
    
    mcp.run(transport="stdio")